import torch
import torch.nn as nn
from models.activator.mish import Mish
# from models.dnn0 import DNN
# from models.dnn_residual import DNN
from models.dnn_residual import DNN
class DNNs(nn.Module):
    def __init__(self, pth_path_list=None):
        super().__init__()
        # self.dnn0 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
        # self.dnn1 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
        # self.dnn2 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
        # self.dnn3 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
        # self.dnn4 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
        self.dnn0 = DNN(28169, num_class=3, hidden_list=[4096, 2048, 1024,1024,1024,512], dropout_p_list=[0.5,0.5,0.4,0.4,0.4,0.2])
        self.dnn1 = DNN(28169, num_class=3, hidden_list=[4096, 2048, 1024,1024,1024,512], dropout_p_list=[0.5,0.5,0.4,0.4,0.4,0.2])
        self.dnn2 = DNN(28169, num_class=3, hidden_list=[4096, 2048, 1024,1024,1024,512], dropout_p_list=[0.5,0.5,0.4,0.4,0.4,0.2])
        self.dnn3 = DNN(28169, num_class=3, hidden_list=[4096, 2048, 1024,1024,1024,512], dropout_p_list=[0.5,0.5,0.4,0.4,0.4,0.2])
        self.dnn4 = DNN(28169, num_class=3, hidden_list=[4096, 2048, 1024,1024,1024,512], dropout_p_list=[0.5,0.5,0.4,0.4,0.4,0.2])
        if pth_path_list != None:
            for i,dnn in enumerate([self.dnn0,self.dnn1,self.dnn2,self.dnn3,self.dnn4]):
                pre_weight = torch.load(pth_path_list[i])
                dnn.load_state_dict(pre_weight)


    def forward(self, x):
        output0=self.dnn0(x)
        output1=self.dnn1(x)
        output2=self.dnn2(x)
        output3=self.dnn3(x)
        output4=self.dnn4(x)
        output = output0+output1+output2+output3+output4
        output = output/5

        return output